#!/usr/bin/perl -w
##This script use BEDtools to annotate bed file with multiple genomic features, with percentage overlap, within same strand, reverse strand or both.
##it extend BEDtools to allow count from 5'end, 3'end, both border or center of bed. it also allows different upstream and downstream. 
##it will generate summary files that count how many regions are overlapped with each feature(p value using hypergenomitric test), 
##how many features are in this region(p value using hypergenomitric test)

## author: Yaping Liu  lyping1986@gmail.com 
## time: 2012-7-25

#Usege:  perl annotateBed.pl [Options] bed/gff_file summary_file.txt --anno feature_bed_file_dir

use strict;
use Getopt::Long;
use File::Basename;

sub usage {

    print "\nUsage:\n";
    print "perl annotateBed.pl [Options] bed/gff_file summary_file.txt --anno feature_bed_file_dir\n\n";

    print " This script use BEDtools to annotate bed file with multiple genomic features, with percentage overlap, within same strand, reverse strand or both.\n";
	print " it extend BEDtools to allow count from 5'end, 3'end, both border or center of bed. it also allows different upstream and downstream. \n";
	print " it will generate summary files that count how many regions are overlapped with each feature(p value using hypergenomitric test), \n";
	print " how many features are in this region(p value using hypergenomitric test)\n\n";
    

	print "  bed/gff_file     Input bed/gff file which requires the annotation.\n";
	print "  summary_file.txt     Summary statistics about annotation result.\n\n";
	
	print "  [Options]:\n\n";
    print "  --anno FILE : genomic features' bed files, which are used to annotate input bed files, allow multiple files\n\n";
    print "  --omit_annotation_step : when specified, it will not do annotation but only do plot step\n\n";
    print "  --upstream NUM :  how many bp upstream from  align start position (Default: 0)\n\n";
    print "  --downstream NUM :   how many bp downstream from  align start position (Default: 0)\n\n";
    print "  --overlap_only : only genomic feature overlapped with .bed region will be taken into account\n\n";
    print "  --overlap_percentage NUM :  minimum percentage of .bed regions overlapped with genomic features will be taken into account (Default: 0)\n\n";
    print "  --same_strand : Require same strandedness.  That is, only counts overlaps on the _same_ strand (Default: both strand).\n\n";
	print "  --reverse_strand : Require different strandedness.  That is, only count overlaps on the _opposite_ strand (Default: both strand).\n\n";
	print "  --bedtools_path : path to BEDtools/bin/.if not specified in PATH environment variable\n\n";
	print "  --r_script FILE :  R script's path and name, which is used for plot\n\n";
	print "  --genome FILE :  genome index file. e.g. hg19.fa.fai file\n\n";
	print "  --align_start NUM : (Default: --align_start 1)\n";
	print "                 1) extend upstream, downstream from both of 5', 3' border\n";
	print "                 2) extend upstream, downstream from 5' border\n";
	print "                 3) extend upstream, downstream from 3' border\n";
	print "                 4) extend upstream, downstream from center of .bed region\n";

    exit(1);
}

##default option setting
my $anno_dir = ();
my $upstream = 0;
my $downstream = 0;
my $overlap_only = "";
my $omit_annotation_step = "";
my $overlap_percentage = 0.0;
my $same_strand = "";
my $reverse_strand = "";
my $align_start = 1;
my $genome = "";
my $r_script = "./annotateBedPlot.R";
my $bedtools_path = "";


GetOptions( "anno_dir=s" => \$anno_dir,
			"omit_annotation_step" => \$omit_annotation_step,
			"upstream=i" => \$upstream,
			"downstream=i" => \$downstream,
			"overlap_only" => \$overlap_only,
			"overlap_percentage=f" => \$overlap_percentage,
			"same_strand" => \$same_strand,
			"reverse_strand" => \$reverse_strand,
			"align_start=i" => \$align_start,
			"genome=s" => \$genome,
			"r_script=s" => \$r_script,
			"bedtools_path=s" => \$bedtools_path);

usage() if ( scalar(@ARGV) == 0 );

if ( scalar(@ARGV) != 2 ) {
    print "Wrong number of arguments\n";
    usage();
}

my $input_bed = $ARGV[0];
my $summary_file = $ARGV[1];

my $tmp_file = $input_bed;

##create new file, if not overlap_only mode
if($overlap_only eq ""){
	if($tmp_file =~ /\.bed/){
		$tmp_file =~ s/\.bed//;
		$tmp_file .= ".up$upstream.down$downstream.bed";
	}
	elsif($tmp_file =~ /\.gff/){
		$tmp_file =~ s/\.gff//;
		$tmp_file .= ".up$upstream.down$downstream.gff";
	}
	elsif($tmp_file =~ /\.gtf/){
		$tmp_file =~ s/\.gtf//;
		$tmp_file .= ".up$upstream.down$downstream.gtf";
	}
	elsif($tmp_file =~ /\.vcf/){
		$tmp_file =~ s/\.vcf//;
		$tmp_file .= ".up$upstream.down$downstream.vcf";
	}
	else{
		die "not support this file type yet! Only allow .bed, .gff, .vcf files\n";
	}
	&creat_location_file($input_bed,$tmp_file);
	
}
else{
	
}
my $anno_out_file = $tmp_file.".annotate_detail.txt";
my $anno_head_file = $tmp_file.".annotate_detail.head.txt";
opendir(DH,"$anno_dir") or die;
my @anno = readdir(DH);

&overlap($tmp_file, $anno_out_file, $anno_head_file);
#&summary_annotation($anno_out_file);
my $wd = dirname($anno_out_file);
my $genome_size = 3101804739;
if($genome ne ""){
	$genome_size = &summary_genome_size($genome);
}
my $R_cmd_plot = "R --no-restore --no-save --args wd=$wd file=$anno_out_file anno_head=$anno_head_file genome_size=$genome_size < $r_script \n";
system($R_cmd_plot);			
closedir(DH);
if($tmp_file ne $input_bed){
	my $cmd = "rm $tmp_file\n";
	system($cmd);
}


sub overlap{
	my $input_bed = shift(@_);
	my $output_bed = shift(@_);
	my $anno_head_file = shift(@_);
	open(OUT,">$anno_head_file") or die;
	my $cmd = "";

	if($bedtools_path ne ""){
		$cmd = $bedtools_path."/";
	}
	$cmd .= "bedtools annotate -names -both "; 
	if($same_strand ne ""){
		$cmd .= "-s ";
	}
	elsif($reverse_strand ne ""){
		$cmd .= "-S ";
	}
	$cmd .= "-i $input_bed -files "; 
	foreach my $ann(@anno){
		next if $ann eq '.' or $ann eq '..';
		
		my $tmp_ann = $anno_dir."$ann";
		$cmd .= "$tmp_ann ";
		my ($size,$count)=&summary_anno_file_size($tmp_ann);
		$ann =~ s/hg19*//;
		print OUT "$ann\t$size\t$count\t";
	}
	print OUT "\n";
	close(OUT);
	$cmd .= " > $output_bed\n";
	print "$cmd\n";
	if($omit_annotation_step eq ""){
		system($cmd);
	}
	

}

sub summary_annotation{
	my $input_bed = shift(@_);
	open(FH,"<$input_bed") or die "can not open file:$!";
	while(<FH>){
		;
	}
	close(FH);
}

sub creat_location_file{
	my $file = shift(@_);
	
	my $location = shift(@_);

	open(FH,"<$file") or die "can not open file:$!";
	open(OUT,">$location") or die "can not open file:$!";
	while(<FH>){
		chomp;
		my $line = $_;
		my @splitin = split "\t",$line;
		if($align_start == 1){
			my $start = $splitin[1] - $upstream;
			my $end = $splitin[2] + $downstream;
			if($#splitin >= 6){
				if($splitin[5] eq "-"){
					$start = $splitin[1] - $downstream;
					$end = $splitin[2] + $upstream;
				}
			}
			$splitin[1] = $start;
			$splitin[2] = $end;
		}
		elsif($align_start == 2 ){
			my $start = $splitin[1] - $upstream;
			my $end = $splitin[1] + $downstream;
			if($#splitin >= 6){
				if($splitin[5] eq "-"){
					$start = $splitin[2] - $downstream;
					$end = $splitin[2] + $upstream;
				}
			}
			$splitin[1] = $start;
			$splitin[2] = $end;
		}
		elsif($align_start == 3 ){
			my $start = $splitin[2] - $upstream;
			my $end = $splitin[2] + $downstream;
			if($#splitin >= 6){
				if($splitin[5] eq "-"){
					$start = $splitin[1] - $downstream;
					$end = $splitin[1] + $upstream;
				}
			}
			$splitin[1] = $start;
			$splitin[2] = $end;
		}
		elsif($align_start == 4 ){
			my $start = int(($splitin[1] + $splitin[2])/2 -1 - $upstream);
			my $end = int(($splitin[1] + $splitin[2])/2 + $downstream);
			if($#splitin >= 6){
				if($splitin[5] eq "-"){
					$start = int(($splitin[1] + $splitin[2])/2 -1 - $downstream);
					$end = int(($splitin[1] + $splitin[2])/2 + $upstream);
				}
			}
			$splitin[1] = $start;
			$splitin[2] = $end;
		}
		else{
			die "Only accept 1,2,3,4 for --align_start option\n";
		}
		my $newLine = join "\t", @splitin;
		print OUT "$newLine\n";
		
	}
	close(FH);
	close(OUT);
}

sub summary_genome_size{
	my $genome = shift @_;
	open(FH,"<$genome") or die;
	my @in = <FH>;
	chomp(@in);
	close(FH);
	my $sum=0;
	foreach my $line(@in){
		my @splitin=split "\t",$line;
		$sum += $splitin[1];
	}
	return $sum;
}

sub summary_anno_file_size{
	my $file = shift @_;
	open(FH,"<$file") or die;
	my @in = <FH>;
	chomp(@in);
	close(FH);
	my $sum=0;
	my $count=0;
	foreach my $line(@in){
		my @splitin=split "\t",$line;
		if($splitin[1] eq $splitin[2]){
			$sum++;
		}
		else{
			$sum += $splitin[2]-$splitin[1];
		}
		$count++;
	}
	return ($sum,$count);
}


